#include "../DWMRICore/HOT.h"
#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"

#include "../DWMRICore/UnitVectors.h"


#include <omp.h>

#include <stdlib.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <string>

using namespace std;

char *pHOTPathname = NULL;			/* input HOT file pathname */
char *pFAPathname = NULL;			/* input FA file pathname */
char *pFiberConfPathname = NULL;		/* output fiber configure file pathname */

float stepSize = 0.1f;
float minLength = 0.0f;
float maxLength = 100.0f;
float angleThreshold = 0.707f;
float faThreshold = 0.1f;
uint fiberNum = 1000;

//float initThreshold = 0.2f;

CHOT *pHOT = NULL;
CScalar *pFA = NULL;

//struct Seed {
//	CPoint3F  pos;
//	CVector3F dir;
//};

void PrintUsage()
{
	printf("Usage: ");
	printf("ProbabilisticFiberTracking -hot <HOT> -fa <FA> -fiber <FiberConf> [-step_size <stepSize> -min_length <minLength> -max_length <maxLength> -fa_threshold <faThreshold> -angle_threshold <angleThreshold>]");
	printf("\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-hot") == 0) {
			pHOTPathname = argv[++i];
		} else if (strcmp(argv[i], "-fa") == 0) {
			pFAPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberConfPathname = argv[++i];
		} else if (strcmp(argv[i], "-step_size") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-min_length") == 0) {
			minLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-max_length") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pHOTPathname == NULL || pFAPathname == NULL || pFiberConfPathname == NULL)
		return -2;

	return 0;
}

bool CheckSeed(Seed &seed, const float threshold)
{
	float hot[28];
	pHOT->GetHOT(seed.pos, hot);

	float v = 0.0f;
	for (uint i = 0; i < pHOT->GetN(); ++i) {
		v += fabs(hot[i]);
	}
	if (v < 10e-7) {
		return false;
	}

	float currValue = pHOT->ComputeValue(hot, seed.dir);

	float maxValue = 0.0f;
	for (int i = 0; i < 50; ++i) {
		CVector3F thisDir;
		thisDir.m_x = (float)rand() / (float)RAND_MAX - 0.5f;
		thisDir.m_y = (float)rand() / (float)RAND_MAX - 0.5f;
		thisDir.m_z = (float)rand() / (float)RAND_MAX - 0.5f;
		thisDir.Normalize();
		float thisValue = pHOT->ComputeValue(hot, thisDir);
		if (thisValue > maxValue)
			maxValue = thisValue;
	}

	maxValue = maxValue * 1.5f;

	if (threshold < currValue/maxValue)
		return true;
	else
		return false;
}

CVector3F NewRandomDirection(CVector3F dir)
{
	float v[3];
	do {
		v[0] = 2.0f * (float)rand() / (float)RAND_MAX - 1.0f;
		v[1] = 2.0f * (float)rand() / (float)RAND_MAX - 1.0f;
	} while (v[0]*v[0] + v[1]*v[1] > 1.0); 

	 v[0] *= 0.2f;
	 v[1] *= 0.2f;
	 v[2] = 1.0f - (v[0]*v[0] + v[1]*v[1]);
	 v[2] = v[2] < 0.0f ? 0.0f : sqrtf(v[2]);

	 if (dir.m_x * dir.m_x + dir.m_y * dir.m_y < 1e-4) {
		 return CVector3F(v[0], v[1], dir.m_z > 0.0f ? v[2] : -v[2]);
	 }

	 CVector3F y = CVector3F(dir.m_x, dir.m_y, 0.0f);
	 y.Normalize();
	 CVector3F x = CVector3F(-y.m_y, y.m_z, 0.0f);
	 CVector3F y2 = CVector3F(-x.m_y*dir.m_z, x.m_x*dir.m_z, x.m_y*dir.m_x-x.m_x*dir.m_y);
	 y2.Normalize();

	float cx = v[0] * x.m_x + v[1] * x.m_y;
	float cy = v[0] * y.m_x + v[1] * y.m_y;

	CVector3F res = CVector3F(
		cx * x.m_x + cy * y2.m_x + v[2] * dir.m_x,
		cx * x.m_y + cy * y2.m_y + v[2] * dir.m_y,
		cx * x.m_z + cy * y2.m_z + v[2] * dir.m_z);
	res.Normalize();
	return res;
}

bool NextDirection(CPoint3F pos, CVector3F lastDir, CVector3F &nextDir)
{
	if (pFA->GetScalar(pos) < faThreshold)
		return false;

	float hot[28];
	memset(hot, 0, sizeof(float)*28);
	pHOT->GetHOT(pos, hot);

	float factor = 0.0f;
	for (int n = 0; n < 81; ++n) {
		CVector3F tempDir(g_vertices[n*3+0], g_vertices[n*3+1], g_vertices[n*3+2]);
		float value = pHOT->ComputeValue(hot, tempDir);
		if (value > factor)
			factor = value;
	}

	float maxVal = 0.0f;
	for (int n = 0; n < 32; ++n) {
		CVector3F tempDir = NewRandomDirection(lastDir);
		float value = pHOT->ComputeValue(hot, tempDir);
		if (value > maxVal)
			maxVal = value;
	}

	if (maxVal < 0.1 * factor) return false;

	maxVal *= 1.5;

	for (int n = 0; n < 64; n++) {
		CVector3F tempDir = NewRandomDirection(lastDir);
		float value = pHOT->ComputeValue(hot, tempDir);

		if (value > 0.1 * factor) {
			if (value > maxVal) {
				printf("max value exceeded! max value = %f, value = %f\n", maxVal, value);
			}
			float randNum = (float)rand() / (float)RAND_MAX;
			if (randNum < value / maxVal) {
				nextDir = tempDir;
				return (true);
			}
		}
	}

	return false;
}

void ProbabilisticFiberTracking()
{
	clock_t time = clock();
	srand(time);

	uint w = pHOT->GetWidth();
	uint h = pHOT->GetHeight();
	uint d = pHOT->GetDepth();

	CFiber *pFibers = new CFiber[fiberNum];
	memset(pFibers, 0, sizeof(CFiber)*fiberNum);
	uint maxSteps = (uint)(maxLength / stepSize);

	for (uint i = 0; i < fiberNum; ++i) {
		/* generate seeds */
		Seed seed;
		do {
			/* position */
			do {
				seed.pos.m_x = (float)rand() / (float)RAND_MAX * (float)w;
				seed.pos.m_y = (float)rand() / (float)RAND_MAX * (float)h;
				seed.pos.m_z = (float)rand() / (float)RAND_MAX * (float)d;

				seed.pos.m_x = 58.0f;
				seed.pos.m_y = 80.0f;
				seed.pos.m_z = 64.0f;
			} while(pFA->GetScalar(seed.pos) < faThreshold);

			/* direction */
			seed.dir.m_x = (float)rand() / (float)RAND_MAX - 0.5f;
			seed.dir.m_y = (float)rand() / (float)RAND_MAX - 0.5f;
			seed.dir.m_z = (float)rand() / (float)RAND_MAX - 0.5f;
			seed.dir.Normalize();
		} while(!CheckSeed(seed, 0.2f));

		/* tracking fiber */
		CPoint3F *pF = new CPoint3F[maxSteps];
		CPoint3F *pB = new CPoint3F[maxSteps];
		uint fnum, bnum;
		CVector3F dir;

		/* forward */
		pF[0] = seed.pos;
		dir = seed.dir;
		for (fnum = 1; fnum < maxSteps; ++fnum) {
			pF[fnum] = pF[fnum-1] + dir * stepSize;
			if (NextDirection(pF[fnum], dir, dir) == false) {
				break;
			}
		}

		/* backward */
		pB[0] = seed.pos;
		dir.m_x = -seed.dir.m_x; dir.m_y = -seed.dir.m_y; dir.m_z = -seed.dir.m_z;
		for (bnum = 1; bnum < maxSteps; ++bnum) {
			pB[bnum] = pB[bnum-1] + dir * stepSize;
			if (NextDirection(pB[bnum], dir, dir) == false) {
				break;
			}
		}

		if (fnum + bnum < 30) {
			i--;
		} else {
			pFibers[i].CreateFiber(pF, pB, fnum, bnum, &seed);
		}
		
		SafeDeleteArray(pF);
		SafeDeleteArray(pB);
	}

	char *confPathname = "C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\Probabilistic_Fibers\\test.conf";
	char *vertPathname = "C:\\Ziang\\Dataset\\HumanConnectome\\792564_processed\\For_Paper\\B3000\\Probabilistic_Fibers\\test.vert";
	SaveFibers(confPathname, vertPathname, pFibers, fiberNum, 1);

	SafeDeleteArray(pFibers);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pHOT = new CHOT();
	pHOT->OpenHOTFile(pHOTPathname);
	pFA = new CScalar();
	pFA->OpenScalarFile(pFAPathname);

	ProbabilisticFiberTracking();

	SafeDelete(pHOT);
	SafeDelete(pFA);

	return 0;
}